#!/usr/bin/env python3
import argparse
import os
import shlex
import signal
import socket
import subprocess
import sys
import time
from datetime import datetime
from pathlib import Path

#   python ./tools/run_all.py --save-dir /home/ubuntu/tlsfuzzer/save --server-cmd "/home/ubuntu/experiments/openssl/apps/openssl s_server -key /home/ubuntu/experiments/openssl/key.pem -cert /home/ubuntu/experiments/openssl/cert.pem  -port 4433  -HTTP" --scripts-dir /home/ubuntu/tlsfuzzer/scripts/


def parse_args():
    p = argparse.ArgumentParser(description="Run all tlsfuzzer scripts with an OpenSSL server.")
    p.add_argument("--save-dir", required=True, help="Directory to save transcripts and logs")
    p.add_argument("--server-cmd", required=True, help="OpenSSL s_server command (quoted)")
    p.add_argument("--scripts-dir", default="scripts", help="Directory containing test scripts (default: scripts)")
    p.add_argument("--python", default=sys.executable, help="Python interpreter to run scripts (default: current)")
    p.add_argument("--host", default="127.0.0.1", help="Server host to wait on (default: 127.0.0.1)")
    p.add_argument("--port", type=int, default=None, help="Server port (if not provided, will try to parse from server-cmd, else 4433)")
    p.add_argument("--per-script-timeout", type=int, default=600, help="Timeout per script seconds (default: 600)")
    return p.parse_args()


def parse_port_from_cmd(cmd: str):
    toks = shlex.split(cmd)
    for i, t in enumerate(toks):
        if t in ("-port", "-accept") and i + 1 < len(toks):
            try:
                return int(toks[i + 1])
            except ValueError:
                pass
        if t.startswith("-port=") or t.startswith("-accept="):
            try:
                return int(t.split("=", 1)[1])
            except ValueError:
                pass
    return None


def wait_for_port(host, port, timeout=20):
    start = time.time()
    while time.time() - start < timeout:
        try:
            with socket.create_connection((host, port), timeout=1.0):
                return True
        except OSError:
            time.sleep(0.2)
    return False


def start_server(cmd, log_file):
    # Start server as its own process group so we can kill cleanly
    with open(log_file, "ab") as lf:
        proc = subprocess.Popen(
            shlex.split(cmd),
            stdout=lf,
            stderr=lf,
            preexec_fn=os.setsid,
            close_fds=True,
        )
    return proc


def stop_server(proc):
    try:
        os.killpg(os.getpgid(proc.pid), signal.SIGTERM)
    except Exception:
        pass
    try:
        proc.wait(timeout=10)
    except Exception:
        try:
            os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
        except Exception:
            pass


def discover_scripts(scripts_dir: Path):
    tests = []
    for p in sorted(scripts_dir.glob("test-*.py")):
        tests.append(p)
    if not tests:
        for p in sorted(scripts_dir.glob("*.py")):
            tests.append(p)
    return tests


def main():
    args = parse_args()
    save_root = Path(args.save_dir).absolute()
    scripts_dir = Path(args.scripts_dir).absolute()
    repo_root = scripts_dir.parent
    save_root.mkdir(parents=True, exist_ok=True)

    port = args.port or parse_port_from_cmd(args.server_cmd) or 4433

    scripts = discover_scripts(scripts_dir)
    if not scripts:
        print(f"No scripts found in {scripts_dir}", file=sys.stderr)
        sys.exit(1)

    summary = []
    session_ts = datetime.utcnow().strftime("%Y%m%d-%H%M%S")
    session_dir = save_root / f"batch-{session_ts}"
    session_dir.mkdir(parents=True, exist_ok=True)

    for script in scripts:
        script_name = script.name
        print(f"==> Running {script_name}")

        run_dir = session_dir / script.stem
        run_dir.mkdir(parents=True, exist_ok=True)

        # Prepare logs
        server_log = run_dir / "server.log"
        script_log = run_dir / "script.log"

        # Start server
        server_proc = start_server(args.server_cmd, server_log)
        if not wait_for_port(args.host, port, timeout=30):
            stop_server(server_proc)
            summary.append((script_name, "SERVER_START_FAIL"))
            print(f"   Server failed to start on {args.host}:{port}")
            continue

        # Set transcript directory for this run; tlsfuzzer will create subfolders per script automatically
        env = os.environ.copy()
        env["TLSFUZZER_TRANSCRIPT_DIR"] = str(run_dir)

        # Run the script
        cmd = [args.python, str(script)]
        with open(script_log, "ab") as lf:
            try:
                res = subprocess.run(
                    cmd,
                    stdout=lf,
                    stderr=lf,
                    env=env,
                    cwd=str(repo_root),
                    timeout=args.per_script_timeout,
                )
                code = res.returncode
            except subprocess.TimeoutExpired:
                code = -999

        # Stop server
        stop_server(server_proc)

        status = "OK" if code == 0 else f"FAIL({code})"
        summary.append((script_name, status))
        print(f"   {script_name} => {status}")

    # Write summary
    with open(session_dir / "SUMMARY.txt", "w") as f:
        for name, status in summary:
            f.write(f"{name}\t{status}\n")

    print("\nAll done. Summary:")
    for name, status in summary:
        print(f" - {name}: {status}")
    print(f"\nArtifacts saved under: {session_dir}")


if __name__ == "__main__":
    main()


